// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <benchmark/benchmark.h>

#include <nebula/core/record_batch.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/testing/random.h>
#include <nebula/types/type.h>
#include <nebula/util/benchmark_util.h>

namespace nebula {

template <typename VT>
static void BatchToTensorSimple(benchmark::State& state) {
  using CType = typename VT::c_type;
  std::shared_ptr<DataType> ty = TypeTraits<VT>::type_singleton();

  const int64_t num_cols = state.range(1);
  const int64_t num_rows = state.range(0) / num_cols / sizeof(CType);
  nebula::random::RandomArrayGenerator gen_{42};

  std::vector<std::shared_ptr<Field>> fields = {};
  std::vector<std::shared_ptr<Array>> columns = {};

  for (int64_t i = 0; i < num_cols; ++i) {
    fields.push_back(field("f" + std::to_string(i), ty));
    columns.push_back(gen_.ArrayOf(ty, num_rows));
  }
  auto schema = std::make_shared<Schema>(std::move(fields));
  auto batch = RecordBatch::create(schema, num_rows, columns);

  for (auto _ : state) {
    ASSERT_OK_AND_ASSIGN(auto tensor, batch->ToTensor());
  }
  state.SetItemsProcessed(state.iterations() * num_rows * num_cols);
  state.SetBytesProcessed(state.iterations() * ty->byte_width() * num_rows * num_cols);
}

void SetArgs(benchmark::internal::Benchmark* bench) {
  for (int64_t size : {kL1Size, kL2Size}) {
    for (int64_t num_columns : {3, 30, 300}) {
      bench->Args({size, num_columns});
      bench->ArgNames({"size", "num_columns"});
    }
  }
}

BENCHMARK_TEMPLATE(BatchToTensorSimple, Int8Type)->Apply(SetArgs);
BENCHMARK_TEMPLATE(BatchToTensorSimple, Int16Type)->Apply(SetArgs);
BENCHMARK_TEMPLATE(BatchToTensorSimple, Int32Type)->Apply(SetArgs);
BENCHMARK_TEMPLATE(BatchToTensorSimple, Int64Type)->Apply(SetArgs);

}  // namespace nebula
